package com.wondersgroup.aida.embeddingdemo;

import com.wondersgroup.aida.tools.ModelUtil;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;

import java.util.List;

/**
 * @author: 紫麒麟vip
 * @create: 2025/2/16
 * Description:
 */
public class VectorSearchDemo {
    public static void main(String[] args) {
        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .baseUrl(ModelUtil.BASE_URI_OPENAI)
                .apiKey(ModelUtil.API_KEY_0PENAI)
                .build();

        Response<Embedding> embed1 = embeddingModel.embed("我的名字叫jszhao");

        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .dimension(1536)
                .build();

        //寻找最近的4条记录
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
        List<EmbeddingMatch<TextSegment>> matchedResult = searchResult.matches();
        for (EmbeddingMatch<TextSegment> embeddingMatch : matchedResult) {
            // 打印相似度
            System.out.println(embeddingMatch.score());
        }

    }
}
